#ifndef TRANSFORM_H
#define TRANSFORM_H

#include <Eigen/Dense>
#include <etl/unordered_map.h>
#include <stdexcept>
#include <vector>


using namespace Eigen;

namespace os::tf {
class FrameNode {
public:
    std::string name;
    FrameNode* parent = nullptr;
    Matrix3d rotation; // 父坐标系到本坐标系的旋转
    Vector3d translation; // 父坐标系到本坐标系的平移
    uint8_t depth = 0; // 节点深度（根节点为0）

    explicit FrameNode(std::string const& name) : name(name) {
    }

    void setTransform(Matrix3d const& R, Vector3d const& t) {
        rotation = R;
        translation = t;
    }
};

class TFSystem {
private:
    // 最多建立32个坐标系
    etl::unordered_map<std::string, FrameNode*, 32> frames_;
    FrameNode* root_ = nullptr;

    // 获取节点到LCA的路径（不包括LCA）
    void getPathToLCA(FrameNode* node, FrameNode* lca, std::vector<FrameNode*>& path) const {
        while (node != lca) {
            path.push_back(node);
            node = node->parent;
        }
    }

    // 优化后的LCA查找算法
    FrameNode* findLCA(FrameNode* a, FrameNode* b) const {
        // 平衡深度
        while (a->depth > b->depth) a = a->parent;
        while (b->depth > a->depth) b = b->parent;

        // 同步上溯
        while (a != b) {
            a = a->parent;
            b = b->parent;
        }
        return a;
    }

public:
    ~TFSystem() {
        for (auto& pair : frames_) delete pair.second;
    }

    void initialize(std::string const& base_frame = "base") {
        if (frames_.empty()) {
            FrameNode* node = new FrameNode(base_frame);
            node->setTransform(Matrix3d::Identity(), Vector3d::Zero());
            root_ = node;
            frames_[base_frame] = node;
        }
    }

    void addFrame(std::string const& new_frame, std::string const& parent_frame, Matrix3d const& R, Vector3d const& t) {
        if (frames_.count(new_frame)) {
            throw std::runtime_error("Frame already exists");
        }

        FrameNode* parent = frames_.at(parent_frame);
        FrameNode* new_node = new FrameNode(new_frame);
        new_node->parent = parent;
        new_node->setTransform(R, t);
        new_node->depth = parent->depth + 1;
        frames_[new_frame] = new_node;
    }

    void updateTransform(std::string const& child_frame, Matrix3d const& R, Vector3d const& t) const {
        FrameNode* node = frames_.at(child_frame);
        node->setTransform(R, t);
    }

    Vector3d transformPoint(Vector3d const& point, std::string const& source_frame, std::string const& target_frame) {
        if (source_frame == target_frame) return point;

        FrameNode* source = frames_.at(source_frame);
        FrameNode* target = frames_.at(target_frame);

        // 查找LCA
        FrameNode* lca = findLCA(source, target);

        // 构建路径
        std::vector<FrameNode*> source_path;
        std::vector<FrameNode*> target_path;
        getPathToLCA(source, lca, source_path);
        getPathToLCA(target, lca, target_path);

        // 从point进行变换
        Vector3d res = point;

        for (auto node : source_path) {
            res = node->rotation * res + node->translation;
        }

        for (auto it = target_path.rbegin(); it != target_path.rend(); ++it) {
            FrameNode* node = *it;
            res = node->rotation.transpose() * (res - node->translation);
        }

        // 最终应用变换
        return res;
    }
};
}
#endif //TRANSFORM_H